const express = require('express');
const { Ollama, OllamaEmbeddings } = require('@langchain/ollama');
const { cosineSimilarity } = require('@langchain/core/utils/math');
const ResultWrapper = require('../../utils/resultWrapper');

const router = express.Router();

// 初始化 Ollama 模型用于文本嵌入
// const ollama = new Ollama({
//   baseUrl: process.env.OLLAMA_BASE_URL || 'http://localhost:11434',
//   model: process.env.OLLAMA_MAIN_MODEL || "deepseek-r1:8b",
//   temperature: 0.7, // 控制输出的随机性
// });
// 初始化 Ollama 嵌入模型
const embeddingsOllama = new OllamaEmbeddings({
  baseUrl: process.env.OLLAMA_BASE_URL || 'http://localhost:11434',
  model: process.env.OLLAMA_MAIN_MODEL || "deepseek-r1:8b"
});

// 文本嵌入接口
router.post('/embed', async (req, res) => {
  try {
    const { text } = req.body;
    
    if (!text) {
      return res.status(400).json({ error: '文本内容不能为空' });
    }

    const embeddings = await embeddingsOllama.embedQuery(text);
    console.log(embeddings);
    res.json(ResultWrapper.success(embeddings));
  } catch (error) {
    res.status(500).json({ error: '服务器内部错误' });
  }
});

// 批量文本嵌入接口
router.post('/embed-batch', async (req, res) => {
  try {
    const { texts } = req.body;
    
    if (!Array.isArray(texts) || texts.length === 0) {
      return res.status(400).json(ResultWrapper.error('文本数组不能为空'));
    }

    const embeddings = await embeddingsOllama.embedDocuments(texts);

    res.json(ResultWrapper.success(embeddings));
  } catch (error) {
    res.status(500).json({ error: '服务器内部错误' });
  }
});

// 文本相似度比较接口
router.post('/similarity', async (req, res) => {
  try {
    const { text1, text2 } = req.body;

    if (!text1 || !text2) {
      return res.status(400).json({ error: '请提供两段待比较的文本' });
    }

    // 获取两段文本的嵌入向量
    const [embedding1, embedding2] = await Promise.all([
      embeddingsOllama.embedQuery(text1),
      embeddingsOllama.embedQuery(text2)
    ]);

    // 计算余弦相似度
    const similarity = cosineSimilarity(embedding1, embedding2);

    res.json({
      similarity,
      text1_embeddings: embedding1,
      text2_embeddings: embedding2
    });
  } catch (error) {
    res.status(500).json({ error: '服务器内部错误' });
  }
});

// 文本聚类接口
router.post('/cluster', async (req, res) => {
  try {
    const { texts, numClusters = 3 } = req.body;

    if (!Array.isArray(texts) || texts.length === 0) {
      return res.status(400).json(ResultWrapper.error('文本数组不能为空'));
    }

    if (texts.length < numClusters) {
      return res.status(400).json({ error: '文本数量必须大于等于聚类数量' });
    }

    // 获取所有文本的嵌入向量
    const embeddings = await embeddingsOllama.embedDocuments(texts);

    // 使用简单的K-means算法进行聚类
    const clusters = new Array(numClusters).fill().map(() => []);
    const centroids = embeddings.slice(0, numClusters); // 简单起见，使用前k个向量作为初始质心

    // 分配文本到最近的聚类
    embeddings.forEach((embedding, index) => {
      let maxSimilarity = -1;
      let clusterIndex = 0;

      centroids.forEach((centroid, i) => {
        const similarity = cosineSimilarity(embedding, centroid);
        if (similarity > maxSimilarity) {
          maxSimilarity = similarity;
          clusterIndex = i;
        }
      });

      clusters[clusterIndex].push({
        text: texts[index],
        embedding,
        similarity: maxSimilarity
      });
    });

    res.json({
      clusters: clusters.map((cluster, index) => ({
        cluster_id: index,
        size: cluster.length,
        texts: cluster.map(item => ({
          text: item.text,
          similarity: item.similarity
        }))
      }))
    });
  } catch (error) {
    res.status(500).json({ error: '服务器内部错误' });
  }
});

module.exports = router;